import baostock as bs
import pandas as pd
import yfinance as yf
from datetime import datetime, timedelta
from typing import Optional, Dict, List, Any, Tuple
import time
import re

from .utils.config import get_settings
from .utils.logger import get_logger

logger = get_logger("AShareFetcher")


class AShareDataFetchError(Exception):
    """A股数据获取异常"""
    pass


class AShareFetcher:
    """A股数据获取器"""

    def __init__(self):
        self.settings = get_settings()
        self.logger = logger
        self.bs = None
        self.is_logged_in = False

    def __del__(self):
        """析构函数，确保登出"""
        if self.is_logged_in:
            self.logout()

    def login(self) -> bool:
        """登录baostock"""
        try:
            if self.is_logged_in:
                return True

            self.logger.info("登录baostock...")
            result = bs.login()
            if result.error_code != '0':
                raise AShareDataFetchError(f"登录失败: {result.error_msg}")

            self.is_logged_in = True
            self.bs = bs
            self.logger.info("baostock登录成功")
            return True

        except Exception as e:
            self.logger.error(f"baostock登录失败: {e}")
            return False

    def logout(self):
        """登出baostock"""
        try:
            if self.is_logged_in and self.bs:
                self.bs.logout()
                self.is_logged_in = False
                self.bs = None
                self.logger.info("baostock已登出")
        except Exception as e:
            self.logger.error(f"baostock登出失败: {e}")

    def _normalize_ashare_code(self, symbol: str) -> Tuple[str, str]:
        """标准化A股代码格式"""
        symbol = symbol.strip().upper()

        # 如果已经是baostock格式，直接返回
        if symbol.startswith(('SH.', 'SZ.', 'BJ.')):
            return symbol, symbol.split('.')[0]

        # 判断交易所并转换为baostock格式
        if symbol.startswith('688'):  # 科创板
            return f"SH.{symbol}", symbol
        elif symbol.startswith('6'):  # 上海主板
            return f"SH.{symbol}", symbol
        elif symbol.startswith('000') or symbol.startswith('001') or symbol.startswith('002'):  # 深圳主板
            return f"SZ.{symbol}", symbol
        elif symbol.startswith('300'):  # 创业板
            return f"SZ.{symbol}", symbol
        elif symbol.startswith('43') or symbol.startswith('83') or symbol.startswith('87'):  # 新三板
            return f"BJ.{symbol}", symbol
        elif symbol.startswith('8'):  # 北交所
            return f"BJ.{symbol}", symbol
        elif symbol.startswith('5') or symbol.startswith('588'):  # 沪市ETF（包括科创板ETF）
            return f"SH.{symbol}", symbol
        elif symbol.startswith('159'):  # 深市ETF
            return f"SZ.{symbol}", symbol
        else:
            # 默认认为是上海证券交易所
            return f"SH.{symbol}", symbol

    def _convert_frequency(self, interval: str) -> str:
        """转换数据间隔格式"""
        frequency_map = {
            '1d': 'd',    # 日线
            '1w': 'w',    # 周线
            '1M': 'm',    # 月线
            '5m': '5',    # 5分钟线
            '15m': '15',  # 15分钟线
            '30m': '30',  # 30分钟线
            '1h': '60',   # 小时线
        }
        return frequency_map.get(interval, 'd')

    def validate_symbol(self, symbol: str) -> bool:
        """验证A股代码格式"""
        if not symbol or not isinstance(symbol, str):
            return False

        symbol = symbol.strip().upper()

        # A股代码格式：6位数字
        if not re.match(r'^\d{6}$', symbol):
            return False

        # 检查代码范围
        valid_ranges = [
            '60',  # 上海主板
            '68',  # 科创板
            '00',  # 深圳主板
            '30',  # 创业板
            '43',  # 新三板
            '83',  # 新三板
            '87',  # 新三板
            '8',   # 北交所
            '51',  # 沪市ETF
            '511', # 沪市ETF
            '512', # 沪市ETF
            '515', # 沪市ETF
            '516', # 沪市ETF
            '517', # 沪市ETF
            '518', # 沪市ETF
            '588', # 科创板ETF
            '159', # 深市ETF
        ]

        return any(symbol.startswith(prefix) for prefix in valid_ranges)

    def get_symbol_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票基本信息"""
        try:
            if not self.validate_symbol(symbol):
                return {'symbol': symbol, 'name': 'Unknown', 'market': 'Unknown'}

            bs_code, clean_code = self._normalize_ashare_code(symbol)

            if not self.login():
                return {'symbol': symbol, 'name': 'Unknown', 'market': 'Unknown'}

            # 获取证券信息
            rs = self.bs.query_stock_basic(bs_code)
            if rs.error_code != '0':
                return {'symbol': symbol, 'name': 'Unknown', 'market': 'Unknown'}

            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if data_list:
                info = data_list[0]
                return {
                    'symbol': symbol,
                    'name': info[1] if len(info) > 1 else 'Unknown',  # 证券名称
                    'market': info[2] if len(info) > 2 else 'Unknown',  # 所属市场
                    'code': info[0] if len(info) > 0 else symbol,     # 证券代码
                    'type': 'A-Share'
                }
            else:
                return {'symbol': symbol, 'name': 'Unknown', 'market': 'Unknown'}

        except Exception as e:
            self.logger.error(f"获取 {symbol} 信息失败: {e}")
            return {'symbol': symbol, 'name': 'Unknown', 'market': 'Unknown'}

    def fetch_history_data(self, symbol: str, interval: str = '1d',
                          start_date: str = None, end_date: str = None) -> pd.DataFrame:
        """
        获取A股历史数据

        Args:
            symbol: 股票代码（如 '688228'）
            interval: 数据间隔（1d, 1w, 1M, 5m, 15m, 30m, 1h）
            start_date: 开始日期（格式：YYYY-MM-DD）
            end_date: 结束日期（格式：YYYY-MM-DD）

        Returns:
            包含历史数据的DataFrame
        """
        try:
            if not self.validate_symbol(symbol):
                raise AShareDataFetchError(f"无效的A股代码: {symbol}")

            bs_code, clean_code = self._normalize_ashare_code(symbol)
            frequency = self._convert_frequency(interval)

            if not self.login():
                raise AShareDataFetchError("baostock登录失败")

            # 设置日期范围
            if not end_date:
                end_date = datetime.now().strftime('%Y-%m-%d')
            if not start_date:
                # 如果是分钟线数据，只获取最近60天
                if interval in ['5m', '15m', '30m', '1h']:
                    start_date = (datetime.now() - timedelta(days=60)).strftime('%Y-%m-%d')
                else:
                    # 根据股票类型确定起始日期
                    if symbol.startswith('688'):  # 科创板
                        start_date = '2019-07-22'  # 科创板开市日期
                    elif symbol.startswith('300'):  # 创业板
                        start_date = '2009-10-30'  # 创业板开市日期
                    elif symbol.startswith('000') or symbol.startswith('001') or symbol.startswith('002'):
                        start_date = '1990-12-19'  # 深圳证券交易所开市日期
                    elif symbol.startswith('6'):
                        start_date = '1990-12-19'  # 上海证券交易所开市日期
                    else:
                        start_date = '2015-01-01'  # baostock数据从2015年开始

            self.logger.info(f"获取 {symbol} ({bs_code}) 数据: {start_date} 到 {end_date}")

            # 定义返回字段
            fields = "date,open,high,low,close,volume,amount,turn,pctChg"

            # 获取数据
            rs = self.bs.query_history_k_data_plus(
                bs_code,
                fields,
                start_date=start_date,
                end_date=end_date,
                frequency=frequency
            )

            if rs.error_code != '0':
                raise AShareDataFetchError(f"数据获取失败: {rs.error_msg}")

            # 读取数据
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                raise AShareDataFetchError("未获取到数据")

            # 创建DataFrame
            df = pd.DataFrame(data_list, columns=rs.fields)

            # 数据清洗和转换
            df = self._clean_dataframe(df, symbol, interval)

            self.logger.info(f"成功获取 {symbol} 数据: {len(df)} 条记录")
            return df

        except Exception as e:
            self.logger.error(f"获取 {symbol} 数据失败: {e}")
            raise AShareDataFetchError(f"获取 {symbol} 数据失败: {e}")

    def _clean_dataframe(self, df: pd.DataFrame, symbol: str, interval: str) -> pd.DataFrame:
        """清洗和转换DataFrame数据"""
        try:
            # 删除空值
            df = df.replace('', None).dropna()

            if df.empty:
                return df

            # 数值列转换
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'amount']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 百分比列转换
            pct_columns = ['pctChg', 'turn']
            for col in pct_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 日期列转换
            if 'date' in df.columns:
                df['date'] = pd.to_datetime(df['date'])

            # 重命名列以匹配yfinance格式
            column_mapping = {
                'date': 'Date',
                'open': 'Open',
                'high': 'High',
                'low': 'Low',
                'close': 'Close',
                'volume': 'Volume',
                'amount': 'Amount',
                'turn': 'Turnover',
                'pctChg': 'Change_pct'
            }

            df = df.rename(columns=column_mapping)

            # 确保Date是列而不是索引（与DataFetcher保持一致）
            if 'Date' in df.columns and df.index.name == 'Date':
                df = df.reset_index()

            # 确保有Date列
            if 'Date' not in df.columns and df.index.name == 'Date':
                df = df.reset_index()

            # 按日期排序
            if 'Date' in df.columns:
                df = df.sort_values('Date')
            else:
                df = df.sort_index()

            return df

        except Exception as e:
            self.logger.error(f"数据清洗失败: {e}")
            return df

    def fetch_incremental_update(self, symbol: str, last_date: str,
                               interval: str = '1d') -> pd.DataFrame:
        """增量更新数据"""
        try:
            # 计算开始日期（最后日期的下一天）
            last_dt = datetime.strptime(last_date, '%Y-%m-%d')
            start_date = (last_dt + timedelta(days=1)).strftime('%Y-%m-%d')
            end_date = datetime.now().strftime('%Y-%m-%d')

            # 如果开始日期就是今天或更晚，没有新数据
            if start_date >= end_date:
                self.logger.info(f"没有新数据需要更新 for {symbol}")
                return pd.DataFrame()

            return self.fetch_history_data(symbol, interval, start_date, end_date)

        except Exception as e:
            self.logger.error(f"增量更新 {symbol} 失败: {e}")
            raise AShareDataFetchError(f"增量更新 {symbol} 失败: {e}")

    def fetch_batch_symbols(self, symbols: List[str], interval: str = '1d') -> Dict[str, pd.DataFrame]:
        """批量获取多个股票数据"""
        results = {}

        for i, symbol in enumerate(symbols):
            try:
                self.logger.info(f"处理 {symbol} ({i+1}/{len(symbols)})")

                if self.validate_symbol(symbol):
                    data = self.fetch_history_data(symbol, interval)
                    results[symbol] = data
                else:
                    self.logger.warning(f"无效的A股代码: {symbol}")
                    results[symbol] = pd.DataFrame()

                # 添加延迟以避免请求过于频繁
                time.sleep(0.5)

            except Exception as e:
                self.logger.error(f"获取 {symbol} 数据失败: {e}")
                results[symbol] = pd.DataFrame()

        success_count = sum(1 for df in results.values() if not df.empty)
        self.logger.info(f"批量获取完成: {success_count}/{len(symbols)} 成功")

        return results

    def get_supported_intervals(self) -> List[str]:
        """获取支持的数据间隔"""
        return ['1d', '1w', '1M', '5m', '15m', '30m', '1h']

    def is_ashare_symbol(self, symbol: str) -> bool:
        """判断是否为A股代码"""
        return self.validate_symbol(symbol)